#
#
import os
import matplotlib.pyplot as plt
from matplotlib import cm
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor


def get_mnist(root: str, train: bool = True) -> MNIST:
    """
    取得 MNIST 数据集.

    Arges
    :root - 数据集根目录
    :train - 是否取得训练集
    """
    root_dir = root
    is_train = train
    is_download = not(os.path.exists(root_dir)) or not os.listdir(root_dir)
    mnist = MNIST(root=root_dir, train=is_train,
                  transform=ToTensor(), download=is_download)
    return mnist


def disp_mnist(mnist: MNIST):
    '''
    展示下载的数据
    '''
    # 数据： torch.Size([60000, 28, 28]) value = [0:255]
    print(mnist.data.size())
    # 标签： torch.Size([60000]) value = [0:9]
    print(mnist.targets.size())
    pass


def mnist_image(mnist: MNIST, idx: int = 0):
    plt.imshow(mnist.data[idx].numpy(), cmap='gray')
    plt.title(f'[{mnist.targets[idx]}] -- vindex = {idx}')
    # plt.title('vindex = {0}, vlabel = {1}'.format(idx, mnist.targets[index]))
    plt.show()
    pass


def plot_with_labels(lowDWeights, labels):
    plt.cla()  # Clear the current axes.
    X, Y = lowDWeights[:, 0], lowDWeights[:, 1]
    for x, y, s in zip(X, Y, labels):
        c = cm.rainbow(int(255 * s / 9))
        plt.text(x, y, s, backgroundcolor=c, fontsize=9)
    plt.xlim(X.min(), X.max())
    plt.ylim(Y.min(), Y.max())
    plt.title('Visualize last layer')
    plt.show()
    plt.pause(0.01)
    pass


